package com.me.study.javaCore.agent;

import com.me.common.util.ExecutingCommand;
import com.me.common.util.PidUtils;
import com.sun.tools.attach.VirtualMachine;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.*;

/**
 * Java Agent 技术（Sun 扩展的非标准技术）：使用 tools.jar 里方法将 agent 动态加载到目标 JVM 的类 <p>
 * 运行时载入，通过 attach API，将模块（jar包）动态地Attach到指定 PID 的 Java进程内     <p>
 *
 * @author ME
 * @date 2021/5/6
 */
public class Attacher {

    private static final Logger LOGGER = LoggerFactory.getLogger(Attacher.class);

    public static void main(String[] args) throws Exception {
        VirtualMachine vm = null;
        try {
            long pid = select(false);
            vm = VirtualMachine.attach(Long.toString(pid)); // 目标 JVM pid
            vm.loadAgent("../testME/testJavaAgent/target/testJavaAgent-my-assembly.jar=thisIsAgentArgs");
        } finally {
            if (null != vm) {
                vm.detach();
            }
        }
    }

    /**
     * 交互式的让用户选择需要 attach的 java进程号（参考 arthas：ProcessUtils.java）
     *
     * @param verbose 是否使用 `jps -v`参数 输出 jvm运行参数
     */
    public static long select(boolean verbose) throws InputMismatchException {
        Map<Long, String> processMap = listProcessByJps(verbose);

        if (processMap.isEmpty()) {
            LOGGER.info("Can not find java process. Try to pass <pid> in command line.");
            return -1;
        }

        LOGGER.info("Found existing java process, please choose one and input the serial number of the process, eg : 1. Then hit ENTER.");
        // print list
        int count = 1;
        for (String process : processMap.values()) {
            if (count == 1) {
                System.out.println("* [" + count + "]: " + process);
            } else {
                System.out.println("  [" + count + "]: " + process);
            }
            count++;
        }

        // read choice
        String line = new Scanner(System.in).nextLine();
        if (line.trim().isEmpty()) {
            // get the first process id
            return processMap.keySet().iterator().next();
        }

        int choice = new Scanner(line).nextInt();

        if (choice <= 0 || choice > processMap.size()) {
            return -1;
        }

        Iterator<Long> idIter = processMap.keySet().iterator();
        for (int i = 1; i <= choice; ++i) {
            if (i == choice) {
                return idIter.next();
            }
            idIter.next();
        }

        return -1;
    }

    /**
     * 借助 jps命令 输出 jvm进程
     *
     * @param verbose 是否使用 `jps -v`参数 输出 jvm运行参数
     */
    private static Map<Long, String> listProcessByJps(boolean verbose) {
        Map<Long, String> result = new LinkedHashMap<Long, String>();

        String jps = "jps";
        File jpsFile = findJps();
        if (jpsFile != null) {
            jps = jpsFile.getAbsolutePath();
        }

        LOGGER.debug("Try use jps to lis java process, jps: " + jps);

        String[] command = null;
        if (verbose) {
            command = new String[] { jps, "-v", "-l" };
        } else {
            command = new String[] { jps, "-l" };
        }

        List<String> lines = ExecutingCommand.runNative(command);

        LOGGER.debug("jps result: " + lines);

        long currentPid = Long.parseLong(PidUtils.currentPid());
        for (String line : lines) {
            String[] strings = line.trim().split("\\s+");
            if (strings.length < 1) {
                continue;
            }
            try {
                long pid = Long.parseLong(strings[0]);
                if (pid == currentPid) {
                    continue;
                }
                if (strings.length >= 2 && isJpsProcess(strings[1])) { // skip jps
                    continue;
                }

                result.put(pid, line);
            } catch (Throwable e) {
                // https://github.com/alibaba/arthas/issues/970
                // ignore
            }
        }

        return result;
    }


    private static boolean isJpsProcess(String mainClassName) {
        return "sun.tools.jps.Jps".equals(mainClassName) || "jdk.jcmd/sun.tools.jps.Jps".equals(mainClassName);
    }

    private static File findJps() {
        // Try to find jps under java.home and System env JAVA_HOME
        String javaHome = System.getProperty("java.home");
        String[] paths = { "bin/jps", "bin/jps.exe", "../bin/jps", "../bin/jps.exe" };

        List<File> jpsList = new ArrayList<File>();
        for (String path : paths) {
            File jpsFile = new File(javaHome, path);
            if (jpsFile.exists()) {
                LOGGER.debug("Found jps: " + jpsFile.getAbsolutePath());
                jpsList.add(jpsFile);
            }
        }

        if (jpsList.isEmpty()) {
            LOGGER.debug("Can not find jps under :" + javaHome);
            String javaHomeEnv = System.getenv("JAVA_HOME");
            LOGGER.debug("Try to find jps under env JAVA_HOME :" + javaHomeEnv);
            for (String path : paths) {
                File jpsFile = new File(javaHomeEnv, path);
                if (jpsFile.exists()) {
                    LOGGER.debug("Found jps: " + jpsFile.getAbsolutePath());
                    jpsList.add(jpsFile);
                }
            }
        }

        if (jpsList.isEmpty()) {
            LOGGER.debug("Can not find jps under current java home: " + javaHome);
            return null;
        }

        // find the shortest path, jre path longer than jdk path
        if (jpsList.size() > 1) {
            Collections.sort(jpsList, new Comparator<File>() {
                @Override
                public int compare(File file1, File file2) {
                    try {
                        return file1.getCanonicalPath().length() - file2.getCanonicalPath().length();
                    } catch (IOException e) {
                        // ignore
                    }
                    return -1;
                }
            });
        }
        return jpsList.get(0);
    }

}
